import torch.nn as nn
import torch
import torch.nn.functional as F

class FDB(nn.Module):
    def __init__(self, cnum):
        super(FDB, self).__init__()
        self.cnum = cnum

    def _CCMP(self,x, kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False):
        x = x.transpose(3, 1)
        x = F.max_pool2d(x, kernel_size, stride,
                         padding, dilation, ceil_mode,
                         return_indices)
        x = x.transpose(3, 1).contiguous()
        return x

    def _SL(self,x):
        n, c, h, w = x.size()
        f = x.view(n, c, h * w)
        index = torch.argmax(x, dim=2)
        b1 = torch.ones(f.size())
        b1 = b1.cuda()
        # 抑制比例
        b1[:, :, index] = 0
        f = f * b1
        f = f.view(n, c, h, w)

        return f

    def forward(self,x, target):
        x = self._CCMP(x, kernel_size=(1, self.cnum), stride=(1, self.cnum))
        x = self._SL(x)
        x = nn.AdaptiveAvgPool2d((1, 1))(x)
        x = x.view(x.size(0), -1)
        out = nn.CrossEntropyLoss()(x, target)

        return out





